import datetime
import threading
import tushare as ts
import pandas as pd
import threading

from sqlalchemy import Column, Integer, String, FLOAT, Boolean, DATE
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import sessionmaker
from sqlalchemy import inspect, UniqueConstraint
from data_resource.data_bases import Base, engine
from contextlib import contextmanager
from concurrent.futures import ThreadPoolExecutor
import numpy as np


class MyThread(threading.Thread):
    def __init__(self, func, args=(), name=None):
        super(MyThread, self).__init__()
        self.result = None
        self.func = func
        self.args = args
        self.name = name or (func.__name__ if func else "Unknown")
        self.exception = None

    def run(self):
        try:
            if self.func is None:
                raise ValueError("Function is None")
            self.result = self.func(*self.args)
        except Exception as e:
            self.exception = e
            print(f"Error in thread '{self.name}': {e}")
            if self.func is not None:
                print(f"Function name: {self.func.__name__}")
            else:
                print("No function assigned")

    def get_result(self):
        try:
            return self.result
        except Exception:
            return None


def init_ts(token_path=r"E:\utilities\quant\tushare_token.txt"):
    with open(token_path, 'r') as f:
        token = f.read()
    ts.set_token(token)
    pro = ts.pro_api()
    return pro


def pandas_dtype_to_sqlalchemy(dtype):
    dtype_map = {
        'int64': Integer,
        'float64': FLOAT,
        'object': String,
        'datetime64[ns]': DATE,
        'bool': Boolean,
        'category': String  # 根据实际需求调整
    }
    return dtype_map.get(str(dtype), String)  # 默认使用 String


# 动态创建ORM类
def create_dynamic_model(table_name, df, engineconn, primary_key='id', unique_columns=None):
    """根据 DataFrame 动态创建 SQLAlchemy 模型类"""
    # 初始化列列表
    columns = [
        Column(primary_key, Integer, primary_key=True, autoincrement=True)
    ]
    table_args = ({'extend_existing': True})

    # 添加索引约束
    if unique_columns:
        for col in unique_columns:
            if col in df.columns:
                columns.append(Column(col, pandas_dtype_to_sqlalchemy(df[col].dtype)))

        constraint_name = f'uq_{table_name}'
        table_args = (
            UniqueConstraint(*unique_columns, name=constraint_name),
            {'extend_existing': True}
        )

    # 添加其他列
    for col in df.columns:
        if unique_columns:
            if col not in [primary_key] + unique_columns:
                columns.append(
                    Column(col, pandas_dtype_to_sqlalchemy(df[col].dtype))
                )
        else:
            if col not in [primary_key]:
                columns.append(
                    Column(col, pandas_dtype_to_sqlalchemy(df[col].dtype))
                )

    # 动态创建模型类
    model = type(
        table_name,  # 类名使用表名的大写形式
        (Base,),
        {
            '__tablename__': table_name,

            '__table_args__': table_args,
            '__mapper_args__': {'always_refresh': True},
            '__module__': __name__,
            **{col.name: col for col in columns}
        }
    )

    # 创建表（如果不存在）
    inspector = inspect(engineconn)
    if inspector.has_table(table_name, schema='quant_research'):
        return model
    else:
        Base.metadata.create_all(engineconn)
        print(f"表{table_name}已创建")
        return model


def bulk_insert_with_orm(df, table_name, engineconn, unique_columns=None, batch_size=1000):
    """使用 ORM 批量插入数据，自动忽略冲突行"""
    # 创建动态模型
    Model = create_dynamic_model(table_name, df, engineconn, unique_columns=unique_columns)

    # 写入数据库
    write_to_db(table_name, engineconn, Model, df)

    return


def clear_financials(df):
    # ts-financials 规整返回数据格式
    df.rename(columns={'ts_code': 'ticker'}, inplace=True)
    df['ticker'] = df['ticker'].apply(lambda x: x[:6])

    df['ann_date'] = pd.to_datetime(df['ann_date'], format='%Y%m%d')
    df['f_ann_date'] = pd.to_datetime(df['f_ann_date'], format='%Y%m%d')
    df['end_date'] = pd.to_datetime(df['end_date'], format='%Y%m%d')
    return df


def clear_forecast(df):
    # ts-financial-forecast 规整返回数据格式
    df.rename(columns={'ts_code': 'ticker'}, inplace=True)
    df['ticker'] = df['ticker'].apply(lambda x: x[:6])

    df['ann_date'] = pd.to_datetime(df['ann_date'], format='%Y%m%d')
    df['end_date'] = pd.to_datetime(df['end_date'], format='%Y%m%d')
    return df


def isEarningSeason(today: datetime.datetime, market_type="Ashare"):
    """
        判断当前日期是否为财报季或预测季
        - A股披露规则
            * 业绩预告：年报预告：不晚于1月31日；一季报预告：4/15前；半年报：7/15前；三季报：10/15前
            * 业绩报告：年报：1-4月完成；一季报：4月；半年报：7-8月；三季报：10月
            * IPO股票的数据容易发生错误（tushare特定）
    """
    if market_type == 'Ashare':
        if today.month == 1:
            return {
                'forecast': True,
                "financials": True,
                "reportdate": [f"{today.year - 1}1231"],
                "forecastdate": f"{today.year - 1}1231"
            }
        elif today.month in [2, 3]:
            return {
                'forecast': False,
                "financials": True,
                "reportdate": [f"{today.year - 1}1231"],
                "forecastdate": None
            }
        elif today.month == 4:
            return {
                'forecast': True,
                "financials": False,
                "reportdate": [f"{today.year - 1}1231", f"{today.year}0131"],  # 去年年报及今年一季报均在4月份披露
                "forecastdate": f"{today.year}0131"
            }
        elif today.month in [6, 7]:
            return {
                'forecast': True,
                "financials": True,
                "reportdate": [f"{today.year}0630", ],
                "forecastdate": f"{today.year}0630"
            }
        elif today.month == 8:
            return {
                'forecast': False,
                "financials": True,
                "reportdate": [f"{today.year}0630", ],
                "forecastdate": None
            }
        elif today.month in [9, 10]:
            return {
                'forecast': True,
                "financials": True,
                "reportdate": [f"{today.year}0930", ],
                "forecastdate": f"{today.year}0930"
            }
        else:
            # 11月，12月不发布财报
            return {
                'forecast': False,
                "financials": False,
                "reportdate": None,
                "forecastdate": None
            }


# 通用数据库写入函数
def _write_single_thread(engineconn, datamodel, data_write: pd.DataFrame):
    error_count = 0

    with get_session(engineconn) as session:

        try:
            session.bulk_insert_mappings(
                datamodel,
                data_write.to_dict('records'),
                return_defaults=False
            )
        except IntegrityError as e:
            session.rollback()
            # print(f"------- 数据插入失败，开始进行逐条插入, 报错{e}--------")
            for _, row in data_write.iterrows():
                try:
                    obj = datamodel(**row.to_dict())
                    session.merge(obj)
                    session.commit()
                except IntegrityError:
                    session.rollback()
                    error_count += 1
        finally:
            session.expunge_all()  # 清楚会话缓存

    return error_count


def _write_multi_thread(engineconn, datamodel, data_write: pd.DataFrame, batch_size=100000, num_threads=20):
    batches = [data_write.iloc[i:i + batch_size] for i in range(0, len(data_write), batch_size)]
    error_count = 0
    lock = threading.Lock()

    def process_batch(batch_data):
        nonlocal error_count
        batch_error = _write_single_thread(engineconn, datamodel, batch_data)

        print(f"线程{threading.current_thread().name}处理完成，错误行数：{batch_error}")

        if batch_error > 0:
            with lock:
                error_count += batch_error

    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        executor.map(process_batch, batches)

    return error_count


def write_to_db(tablename, engineconn, datamodel, data_write: pd.DataFrame, batch_size=100000,
                threshold=10000, num_threads=8):
    """
        分批写入，忽略重复数据，支持大数据多线程写入
    """

    total_rows = len(data_write)
    if total_rows < threshold:
        print(f"********* 数据总量：{total_rows}, 小于{threshold}, 采用单线程写入 **********")
        error_count = _write_single_thread(engineconn, datamodel, data_write)
    else:
        print(f"********* 数据总量：{total_rows}，大于{threshold}, 采用多线程写入，线程数量:{num_threads} **********")
        error_count = _write_multi_thread(engineconn, datamodel, data_write,
                                          batch_size, num_threads)

    print(f"\n 表{tablename}:   数据导入完成：共 {total_rows} 行，失败 {error_count} 行")
    return


# 通用数据库查询资源
@contextmanager
def get_session(engineconn):
    Session = sessionmaker(bind=engineconn)
    session = Session()
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()


# dataframe滑动窗口获取多列数据
def sliding_window(df, window_size, output='array'):
    """
    将 DataFrame 转换为滑动窗口结构：返回 3D array
    :param df: 原始 DataFrame（行数 >= window_size）
    :param window_size: 滑动窗口大小
    :param output: 控制滑动窗口格式，array / dataframe
    :return: np.ndarray, shape = (n_windows, window_size, n_features)
    """
    values = df.values
    n_samples = len(df) - window_size + 1

    if output == 'array':
        try:
            _r = np.lib.stride_tricks.sliding_window_view(
                values, (window_size, values.shape[1])
            )[:, 0, :, :]  # 输出维度为：(n_windows, window_size, n_features)

            # _r = np.stack([values[i:i + window_size] for i in range(n_samples)])
            return _r
        except Exception as e:
            print(e)
            raise ValueError(f"滑动窗口生成失败, 报错信息：{e}")

    elif output == 'dataframe':
        try:
            _r = [df.iloc[i:i + window_size, :] for i in range(n_samples)]
            return _r
        except Exception as e:
            print(e)
            raise ValueError(f"滑动窗口生成失败, 报错信息：{e}")


def get_next_tradeDate(tradeDate: datetime.date, tradings: pd.DataFrame, n: int = 1):
    """
        根据给定日期、交易日序列，返回后n个交易日日期
    """
    required_columns = ['cal_date', 'is_open']
    missing_columns = [col for col in required_columns if col not in tradings.columns]
    assert not missing_columns, f"缺少必要的列: {missing_columns}"

    if n > 0:
        # 返回后交易日
        tradings.sort_values(by='cal_date', ascending=True, inplace=True)
        _next = tradings[(tradings['cal_date'] > tradeDate) & (tradings['is_open'] == 1)]  # 保证输出结果为交易日
        _next.reset_index(drop=True, inplace=True)
        return _next.loc[n - 1, 'cal_date']
    elif n < 0:
        # n为负值时，返回前交易日
        tradings.sort_values(by='cal_date', ascending=False, inplace=True)
        _next = tradings[(tradings['cal_date'] < tradeDate) & (tradings['is_open'] == 1)]
        _next.reset_index(drop=True, inplace=True)
        return _next.loc[abs(n) - 1, 'cal_date']
    else:
        # n=0, 当前是否为交易日
        _temp = tradings[tradings['cal_date'] == tradeDate]['is_open'].values[0]
        if _temp == 1:
            return True
        else:
            return False


def get_price(code: list, startdate: str = '2014-03-31', enddate: str = '2020-12-31', freq='d'):
    """
        获取行情数据
    """

    assert freq in ['d', ]
    _sql = ""

    if freq == 'd':
        if len(code) == 1:
            # 当只有一个代码时，tuple格式化会变成 (code,) 的形式，需要特殊处理
            code_tuple = f"('{code[0]}')"
        else:
            code_tuple = tuple(code)

        _sql = f"""
            select ticker as code, trade_date, open, close, pct_chg, vol 
            from quant_research.market_daily_ts 
            where ticker in {code_tuple} and trade_date between '{startdate}' and '{enddate}'
        """

    price = pd.read_sql(_sql, con=engine)
    price['trade_date'] = pd.to_datetime(price['trade_date']).dt.date
    return price


# 设置涨跌停板类型
def set_limit_type(code: str):
    if code[0] == '3' or code[0:3] == '688':
        return 1  # 20%停板限制
    else:
        return 0  # 10%停板限制


if __name__ == '__main__':
    a = isEarningSeason(datetime.datetime.today())